#
# Created By: Kevin T. Ryan
# Created On: 2007-Feb-28 @ 3:30:45 PM
# Email:      kevin.t.ryan@gmail.com
# License:    BSD (http://www.opensource.org/licenses/bsd-license.php)
#
# New graph module, written from complete scratch.  *Very* beta, assume very many bugs @
# your own risk!
#

import gd       # Only req'd external dependency
import defaults # configuration file specifying what to do for certain instances

class Point(object):

    def __init__(self, x, y):
        self.x = x
        self.y = y

    def pts(self):
        return (self.x, self.y)

    def __str__(self):
        return "<Point x=%d y=%d>" % (self.x, self.y)

    def __repr__(self):
        return self.__str__()

class Graph(object):

    def __init__(self, width=defaults.graph_width, height=defaults.graph_height, xdata=None, ydata=None, **props):
        '''Instantiate the graph image.'''
        self.img = gd.image((width, height))
        # We also set the background before doing anything.  By default, we assume a white
        # background but it's easy enough to overwrite by setting the "background" in the
        # keyword argument list.
        self.img.colorAllocate(props.get("background", defaults.background))
        # We want to keep a list of xvals / yvals as opposed to only one data
        # set ... makes it more robust
        if xdata is not None and type(xdata[0]) != list:
            # ie, data provided but only one dataset ... so we massage it a little bit
            self.xdata = [xdata]
        else:
            self.xdata = []

        if ydata is not None and type(ydata[0]) != list:
            # same as above
            self.ydata = [ydata]
        else:
            self.ydata = []

        self.props = {}
        for k,v in props.items():
            self.props[k] = v

        self._determine_origin()
        # Now that we've deterimined the origin and extent, we can create our scaling
        # function (ie, to scale individual points when graphing).
        self.set_scaling_fn()

    def generate_scaling_fn(self):
        '''Provides a scaling factor (a function) to be used to graph datapoints.

    See the 'SCALING' section in the README.txt for more details.
        '''
        xdatapoints = self.xdata
        ydatapoints = self.ydata
        origin = self.origin
        xextent = self.x_extent
        yextent = self.y_extent
        if not xdatapoints or not ydatapoints:
            return # Can't create a scaling function with no data
        xMin = min([min(vals) for vals in xdatapoints])
        yMin = min([min(vals) for vals in ydatapoints])
        xexpanse = max([max(vals) for vals in xdatapoints]) - xMin
        yexpanse = max([max(vals) for vals in ydatapoints]) - yMin
        xscale = float(xextent.x - origin.x) / xexpanse
        yscale = float(yextent.y - origin.y) / yexpanse

        def _fn(point):
            yscaled = origin.y + (point.y - yMin) * yscale
            xscaled = origin.x + (point.x - xMin) * xscale
            return Point(xscaled, yscaled)
        return _fn

    def _determine_origin(self):
        '''Determines the origin and the extent of our graph for both the x and y axises.'''
        ls = self.get_prop("left_space", defaults.left_space)
        bs = self.get_prop("bottom_space", defaults.bottom_space)
        rs = self.get_prop("right_space", defaults.right_space)
        ts = self.get_prop("top_space", defaults.top_space)

        w,h = self.img.size()
        self.origin   = Point(ls, bs)
        self.x_extent = Point(w - rs, bs)
        self.y_extent = Point(w - rs, h - ts)
        return

    def set_scaling_fn(self):
        self.scale = self.generate_scaling_fn()
        return

    def get_prop(self, prop, default=None):
        return self.props.get(prop, default)

    def set_prop(self, prop, newvalue):
        self.props[prop] = newvalue
        return True

    def _get_font_props(self, props, default_prop_col, default_prop_size):
        '''Returns the font color specified in props, or the default font color from the config file.'''
        fontcolor = props.get("fontcolor", None)
        fontsize  = props.get("fontsize", None)
        if fontcolor is None:
            fontcolor = self.get_prop(default_prop_col, defaults.__dict__[default_prop_col])
        if fontsize is None:
            fontsize = self.get_prop(default_prop_size, defaults.__dict__[default_prop_size])

        if type(fontcolor) != int:
            fontcolor = self.img.colorAllocate(fontcolor)
        return fontcolor, fontsize

    def _get_label_placement(self, fontsize, text):
        '''Returns where the text should be placed on the x or y axis (in that order) to be centered across the image.'''
        # This function will return the width/height of the image by
        # calculating how much space this particular text will take up on this
        # graph ... so to get it centered, we take the width/height of the
        # graph (ie, w/h depending on what axis you want to write on) and subtract
        # the width of the string (ie, sw) and divide by 2 to get the left
        # starting point
        w,h = self.img.size()
        sw,sh = gd.fontstrsize(fontsize, text)
        x_placement = int((w - sw) / 2.0)
        y_placement = int((h - sw) / 2.0)
        return x_placement, y_placement

    def _get_label_text(self, text, default_text_location):
        '''Returns either the text that the user supplied or one specified in the config file/properties of this instance.'''
        if not text:
            text = self.get_prop(default_text_location)
        if not text:
            return None
        return text

    def add_title(self, text=None, **props):
        '''Add's a title to the graph with the associated props (eg, fontsize, fontcolor, etc.).'''
        fontcolor, fontsize = self._get_font_props(props, "titlefontcolor", "titlefontsize")
        text = self._get_label_text(text, "title")
        if not text:
            return # Nothing to type, nothing to do :)

        # Now that we've set up the color, the size and the text, the last step
        # is to simply place the title on the graph
        xpxl_placement, ignore = self._get_label_placement(fontsize, text)
        ypxl_placement = 10 # Always start 10pxls from the top ...
        self.img.string(fontsize, (xpxl_placement, ypxl_placement), text, fontcolor)
        return True

    def add_x_label(self, text=None, **props):
        '''Add's a label to the x-axis so that we know what the numbers across that axis represent.'''
        fontcolor, fontsize = self._get_font_props(props, "xaxisfontcolor", "xaxisfontsize")
        text = self._get_label_text(text, "xlabel")
        if not text:
            return

        w,h = self.img.size()
        sw,sh = gd.fontstrsize(fontsize, text)
        xpxl_placement, ignore = self._get_label_placement(fontsize, text)
        ypxl_placement = h - sh - 6
        self.img.string(fontsize, (xpxl_placement, ypxl_placement), text, fontcolor)
        return True

    def add_y_label(self, text=None, **props):
        fontcolor, fontsize = self._get_font_props(props, "yaxisfontcolor", "yaxisfontsize")
        text = self._get_label_text(text, "ylabel")
        if not text:
            return

        w,h = self.img.size()
        xpxl_placement = 6
        ignore, ypxl_placement = self._get_label_placement(fontsize, text)
        ypxl_placement = h - ypxl_placement
        self.img.stringUp(fontsize, (xpxl_placement, ypxl_placement), text, fontcolor)
        return

    def add_axises(self, color=defaults.axis_color):
        '''Add's an x and a y axis to our graph.'''
        # If the color hasn't been allocated, we'll do so here.  By default, we've chosen
        # a dark gray line to represent the axises
        if type(color) != int:
            color = self.img.colorAllocate(color)
        origin = Point(*self.origin.pts())

        # Set it to the height for drawing purposes (coming from the top)
        w,h = self.img.size()
        origin.y = h - origin.y

        # The x point for the y-axis doesn't change from the origin's x point (ie, it goes
        # straight up).
        upto   = Point(origin.x, h - self.y_extent.y)
        overto = Point(self.x_extent.x, origin.y)

        # Now, draw both lines ... they both must start at the origin, and go over until
        # we've hit our limit
        self.img.line(origin.pts(), upto.pts(), color)
        self.img.line(origin.pts(), overto.pts(), color)

        # Now we add a little extra at the ends to give the graph a little flavor
        downextra = Point(*origin.pts())
        downextra.y += 5
        upextra = Point(*origin.pts())
        upextra.x -= 5
        self.img.line(origin.pts(), downextra.pts(), color)
        self.img.line(origin.pts(), upextra.pts(), color)
        return

    def _get_grid_props(self, props, which_axis):
        '''"Private" method that returns a set of properties to be used when drawing our grid.'''
        color    = props.get('gridcolor', defaults.gridcolor)
        spacing  = props.get('spacing', defaults.gridspacing)
        linelen  = props.get('linelen', defaults.gridstyle)
        # Label specific choices
        labelfnt = props.get('labelfont', defaults.labelfont)
        labelcol = props.get('labelcolor', defaults.labelcolor)

        # The number of lines can be independently set by axis
        if which_axis == "x":
            numlines = float(props.get('numxgridlines', defaults.numxgridlines))
            labels_only = not props.get('showxgrid', defaults.showxgrid)
        elif which_axis == "y":
            numlines = float(props.get('numygridlines', defaults.numygridlines))
            labels_only = not props.get('showygrid', defaults.showygrid)

        # Now we have to transform some of the variables we got to conform with
        # values that gd expects
        if type(color) != int:
            color = self.img.colorAllocate(color)
        if type(labelcol) != int:
            labelcol = self.img.colorAllocate(labelcol)

        return color, spacing, linelen, numlines, labelfnt, labelcol, labels_only

    def draw_x_grids(self, **props):
        '''Adds a grid going up from x-axis of the graph towards the title (ie, vertically).'''
        if self.get_prop("xticklabelsoff"):
            return
        if not props:
            props = self.props

        # Here we just get a list of properties to be used when we're drawing
        # the grid ... see the method definition for how each particular value
        # is obtained.  We also spell out the variables in the 'draw_y_grids'
        # method which is precisely the same as this method, albeit for the y
        # values instead.
        c, sp, ll, nl, lf, lc, lo = self._get_grid_props(props, "x")

        # We don't need to calculate the size of the x text, b/c we go below
        # the grid and hence don't have to worry how many pixels to move the
        # grid down, b/c the text goes down by default
        w,h = self.img.size()
        ylabel_pxl = h - self.origin.y + 5

        # Also, add a label to the x-origin
        xval = self.interpolate_x(self.origin.x)
        xlab = props.get("xlabeldisplay", defaults.xlabeldisplay)(xval)
        self.img.string(lf, (self.origin.x - 7, ylabel_pxl), xlab, lc)

        # How much to step across the grid by
        x_dist = (self.x_extent.x - self.origin.x) / nl
        x_dist = int(x_dist) - 1
        # For each xtick going across the x-axis upto the extent of x, stepping
        # by the distance needed to generate nl (number of lines)
        for xtick in range(self.origin.x + x_dist, self.x_extent.x, x_dist):
            # Interpolate the value of the current xtick and write it to the x-axis
            xval = self.interpolate_x(xtick)
            xlab = props.get("xlabeldisplay", defaults.xlabeldisplay)(xval)
            self.img.string(lf, (xtick - 7, ylabel_pxl), xlab, lc)
            # Also, draw a little line from the x-axis down towards the label
            # for graphical enhancement
            self.img.line((xtick, h - self.origin.y), (xtick, h - self.origin.y + 4), lc)
            if not lo: # if not labels only, draw the grid lines
                # Not only labels, so go from the true origin of y on the image
                # upto the true extent of y on the image stepping by the
                # spacing plus the line length
                for ytick in range(h - self.origin.y, h - self.y_extent.y, -(sp + ll)):
                    endpt = (xtick, ytick + ll)
                    self.img.line((xtick, ytick), endpt, c)
        return True

    def draw_y_grids(self, **props):
        '''Adds a grid going from the y-axis of the graph towards the right edge of the graph (ie, horizontally).'''
        # Get either the props passed to the function, or the one's stored in this instance
        if self.get_prop("yticklabelsoff"):
            return
        if not props:
            props = self.props
        # And set up some values that we need throughout this procedure
        color    = props.get('gridcolor', defaults.gridcolor)
        spacing  = props.get('spacing', defaults.gridspacing)
        linelen  = props.get('linelen', defaults.gridstyle)
        numlines = float(props.get('numlines', defaults.numygridlines))
        labelfnt = props.get('labelfont', defaults.labelfont)
        labelcol = props.get('labelcolor', defaults.labelcolor)

        labels_only = not props.get('showygrid', defaults.showygrid)

        if type(color) != int:
            color = self.img.colorAllocate(color)
        if type(labelcol) != int:
            labelcol = self.img.colorAllocate(labelcol)

        max_y_len = max([max(vals) for vals in self.ydata])
        sw,sh = gd.fontstrsize(labelfnt, props.get("ylabeldisplay", defaults.ylabeldisplay)(max_y_len))
        xlabel_pxl = self.origin.x - (sw + 5)

        # go from the origin + 1 distance mark (ie, don't go across the x axis
        # line) to the extent of y stepping it up by y_dist each time ... this
        # should give us 5 grid lines (ie, across the major ticks)
        y_dist = (self.y_extent.y - self.origin.y) / numlines
        w,h = self.img.size()
        y_dist = int(y_dist) - 1
        # Add a label to the y-origin
        yval = self.interpolate_y(self.origin.y)
        ylab = props.get("ylabeldisplay", defaults.ylabeldisplay)(yval)
        self.img.string(labelfnt, (xlabel_pxl, h - self.origin.y - 7), ylab, labelcol)
        for y in range(self.origin.y + y_dist, self.y_extent.y, y_dist):
            ytick = h - y
            yval  = self.interpolate_y(y)
            ylab  = props.get("ylabeldisplay", defaults.ylabeldisplay)(yval)
            self.img.string(labelfnt, (xlabel_pxl, ytick - 10), ylab, labelcol)
            # Also, add little ticks next to the label for graphical appearances :)
            self.img.line((self.origin.x, ytick), (self.origin.x - 3, ytick), labelcol)
            if labels_only:
                continue
            for xtick in range(self.origin.x, self.x_extent.x, spacing + linelen):
                endpt = (xtick + linelen, ytick)
                self.img.line((xtick, ytick), endpt, color)
        return True

    def save_to_file(self, fn="cache.png"):
        if type(fn) == str:
            t = fn.split(".")
            if len(t) == 1:
                t = defaults.graph_output
            else:
                t = t[-1]
        else:
            t = defaults.graph_output
        fnmap = {'png' : self.img.writePng, 'jpeg' : self.img.writeJpeg}
        func = fnmap.get(t, self.img.writePng)
        func(fn)
        return True

    def clear(self):
        self.xdata = []
        self.ydata = []
        self.get_prop("colors", defaults.colors).reset()
        self.get_prop("barcolors", defaults.barcolors).reset()
        return True

    def add_x_data(self, data):
        '''Set's our xdata values/points.  Also updates our scaling function.'''
        self.xdata.append(data)
        self.set_scaling_fn()

    def add_y_data(self, data):
        '''Set's our ydata values/points.  Also updates our scaling function.'''
        self.ydata.append(data)
        self.set_scaling_fn()

    def interpolate_x(self, pt):
        # See note for interpolate_y function ... contains add'l info on this
        # procedure
        pxl_dist  = self.x_extent.x - self.origin.x
        xMin = min([min(vals) for vals in self.xdata])
        xMax = max([max(vals) for vals in self.xdata])
        val_dist  =  xMax - xMin
        pt_factor = float(pt - self.origin.x) / pxl_dist

        return pt_factor * val_dist + xMin

    def interpolate_y(self, pt):
        # Remember, our ydata variable is a list of lists ... ie, it can
        # contain more than one dataset.  So we take the min/max of each list
        # within the ydata values and then take the min/max of those values
        pxl_dist  = self.y_extent.y - self.origin.y
        yMin = min([min(vals) for vals in self.ydata])
        yMax = max([max(vals) for vals in self.ydata])
        val_dist  = yMax - yMin
        pt_factor = float(pt - self.origin.y) / pxl_dist

        return pt_factor * val_dist + yMin

    def render(self, props=None, fn=None, with_labels=False, add_axises=True):
        '''Default renderer does nothing other than add axises, lables, grids, etc. to the graph.'''
        self.img.setThickness(1)
        # By default, we give labels to both the x and y axises ... you can
        # optionally turn this feature off by setting the property
        # "xticklabelsoff" or "yticklabelsoff" when you create the graph.
        # Also, you can change in the config file whether or not grids are
        # drawn across the 2 axises as well.
        self.draw_y_grids()
        self.draw_x_grids()
        if add_axises:
            self.add_axises()
        if self.get_prop("title") is not None:
            self.add_title()
        if self.get_prop("xlabel") is not None:
            self.add_x_label()
        if self.get_prop("ylabel") is not None:
            self.add_y_label()


class LineGraph(Graph):
    '''Extends graph to draw line graphs.'''

    def render(self, props=None, fn=None, with_labels=False, add_axises=True):
        if props is None:
            props = self.props
        colors = props.get("colors", defaults.colors)
        # Set up the grids / axises / etc.
        Graph.render(self, props, fn, with_labels)

        if props.get("linewidth", defaults.linewidth) != 1:
            self.img.setThickness(props.get("linewidth", defaults.linewidth))
        # We need to scale the points so that they're appropriately placed on the graph.
        # We use our "scale" function, which is dynamically defined at the creation of our
        # graph or at the request of the user.
        points = []
        for xvals,yvals in zip(self.xdata, self.ydata):
            points.append([self.scale(Point(*pt)) for pt in zip(xvals, yvals)])

        # We have to invert our y data ... unfortunately, the gd module makes the top
        # left (instead of the bottom left) the origin of the image.
        w,h = self.img.size()
        for vals in points:
            for pt in vals:
                pt.y = int(h - pt.y)
                pt.x = int(pt.x)
        for j,vals in enumerate(points):
            # Each j represents a new dataset provided by the user.  For each
            # dataset, we grab the next color from our color palette and check
            # if that color has already been allocated, or whether we need to
            # allocate a new color.  By default, when you allocate a color on
            # an image, the gd package will give you back an integer
            # representing the spot where that color was allocated.  Therefore,
            # if the color variable is now an integer we assume that it has
            # been allocated on the image.  If it is a list, then we allocate
            # it ourselves and update the color variable to the integer value
            # returned from the gd package (we also update our colorpalette so
            # we don't have to reallocate again).
            color = colors.next()
            if type(color) != int:
                color = self.img.colorAllocate(color)
                colors.update_last_index(color)
            for i,pt in enumerate(vals[1:]):
                # The previous point is actually equal to points[i] b/c we are enumerating
                # over the points list starting at position 1 (ie, skipping the first value).
                prevPt = vals[i]
                # We also convert each point.x and point.y to an integer; our scaling function
                # can create floats (in fact, *does* create floats by default), but the gd
                # module wants them to be ints (deprecation warning) ... so we'll oblige.
                self.img.line((prevPt.x, prevPt.y), (pt.x, pt.y), color)
                if with_labels:
                    x, y = (self.xdata[j][i], self.ydata[j][i])
                    font = self.get_prop("labelfont", defaults.labelfont)
                    fcolor = self.get_prop("labelcolor", color)
                    self.img.string(font, (prevPt.x, prevPt.y), "<%d, %d>" % (x,y), fcolor)

        if fn is not None:
            self.save_to_file(fn)

class BarGraph(Graph):

    def generate_scaling_fn(self):
        '''Provides a scaling factor (a function) to be used to graph datapoints.

    See the 'SCALING' section in the README.txt for more details.
        '''
        ydatapoints = self.ydata
        origin = self.origin
        yextent = self.y_extent
        if not ydatapoints:
            return # Can't create a scaling function with no data
        yMin = min([min(vals) for vals in ydatapoints], 0)
        yexpanse = max([max(vals) for vals in ydatapoints]) - yMin
        yscale = float(yextent.y - origin.y) / yexpanse

        def _fn(point):
            yscaled = origin.y + (point.y - yMin) * yscale
            return Point(-1, yscaled)
        return _fn

    def interpolate_y(self, pt):
        pxl_dist  = self.y_extent.y - self.origin.y
        yMin = min([min(vals) for vals in self.ydata], 0)
        yMax = max([max(vals) for vals in self.ydata])
        val_dist  = yMax - yMin
        pt_factor = float(pt - self.origin.y) / pxl_dist

        return pt_factor * val_dist + yMin

    def calculate_bar_width(self, origingap, bargap):
        '''Calculates the appropriate width so that we can show all datasets.'''
        num_datasets = len(self.ydata)
        num_datapnts = len(self.ydata[0])
        num_bars = num_datasets * num_datapnts # How many bars will be on the graph
        width_available = (self.x_extent.x - self.origin.x) - origingap - (num_datapnts - 1) * bargap
        barwidth = round(width_available / float(num_bars), 0)
        return int(barwidth)

    def render(self, props=None, fn=None, with_labels=False, add_axises=True):
        '''Extends Graph and implements bar graphs.'''
        Graph.render(self, props, fn, with_labels, add_axises)
        if props is None:
            props = self.props
        colors = props.get("barcolors", defaults.barcolors)

        # The following widths are used to place bars across the x-axis
        bargap = props.get("bargap", defaults.bargap)
        origingap = props.get("origingap", defaults.origingap)
        barwidth = props.get("barwidth", defaults.barwidth)
        if barwidth is None: # Calculate it ourselves
            barwidth = self.calculate_bar_width(origingap, bargap)
            totalwidth = barwidth * len(self.ydata)

        w,h = self.img.size()
        for j,dataset in enumerate(self.ydata):
            color = colors.next()
            if type(color) != int:
                color = self.img.colorAllocate([c for c in color])
                colors.update_last_index(color)

            for i,datapnt in enumerate(dataset):
                uppercorner = self.scale(Point(j, datapnt))
                # Remember, always invert the y values :(
                uppercorner.y = int(h - uppercorner.y)
                # To calculate where the x value should be, we use the formula
                # alluded to in the README.  Note that we have to use j * i b/c
                # if we have multiple datasets, using just the simple formula
                # in the readme will overwrite work we've already done.
                uppercorner.x = int(self.origin.x + origingap + (i * totalwidth) + (i * bargap) + (j * barwidth))
                lowercorner = Point(uppercorner.x + barwidth, h - self.origin.y)
                # TODO: Change the "0" below (which simply puts a border with
                # the background color) to be customizable by the user!
                self.img.rectangle(uppercorner.pts(), lowercorner.pts(), 0, color)

        # We have to draw our grids/axises again to "overwrite" the bars ...
        # otherwise, they overlap the grid and it doesn't look very appealing
        if add_axises:
            self.add_axises()
        if fn is not None:
            self.save_to_file(fn)

    def draw_x_grids(self, **props):
        if self.get_prop("xticklabelsoff"):
            return
        if not props:
            props = self.props

        # Get the x-axis properties (see the main Graph class for more detailed comments)
        c, sp, ll, nl, lf, lc, lo = self._get_grid_props(props, "x")
        w,h = self.img.size()
        ylabel_pxl = h - self.origin.y + 5

        x_dist = (self.x_extent.x - self.origin.x) / nl
        x_dist = int(x_dist) - 1
        # First, we'll do the labels
        nextpt_factor = (self.x_extent.x - self.origin.x) / float(len(self.ydata[0]))
        for i,label in enumerate(self.xdata[0]):
            try:
                label = props.get("xlabeldisplay", defaults.ylabeldisplay)(label)
            except TypeError:
                label = str(label)
            sw,sh = gd.fontstrsize(lf, label)
            start = int(nextpt_factor * i + props.get("origingap", defaults.origingap) + self.origin.x)
            if sw >= start + nextpt_factor:
                self.img.string(lf, (start, ylabel_pxl), label, lc)
            else:
                ctrd = int((nextpt_factor - sw) / 2) + start
                self.img.string(lf, (ctrd, ylabel_pxl), label, lc)

        # Now we add gridlines if the user wants us to
        if not lo:
            for xtick in range(self.origin.x + x_dist, self.x_extent.x, x_dist):
                for ytick in range(h - self.origin.y, h - self.y_extent.y, -(sp + ll)):
                    endpt = (xtick, ytick + ll)
                    self.img.line((xtick, ytick), endpt, c)

        return True

# Some convenience methods
BarChart = BarGraph
LineChart = LineGraph
